#!/usr/bin/env python
# -*- coding: utf-8 -*-

"""
马尔可夫链模型实现
用于彩票预测的状态转移分析
"""

import numpy as np
import pandas as pd
import torch
import os
from collections import defaultdict

class PLWMarkovChainModel:
    """排列5马尔可夫链模型"""
    
    def __init__(self, csv_file_path=None):
        """初始化马尔可夫链模型"""
        self.csv_file_path = csv_file_path
        self.transition_matrices = {}  # 每个位置的状态转移矩阵
        self.initial_probabilities = {}  # 每个位置的初始概率分布
        self.position_stats = {}  # 位置统计信息
        self.is_trained = False
        
    def load_history_data(self, csv_file_path=None):
        """
        加载历史数据
        
        Args:
            csv_file_path: CSV文件路径
            
        Returns:
            历史数据数组
        """
        if csv_file_path is None:
            csv_file_path = self.csv_file_path
            
        if not csv_file_path or not os.path.exists(csv_file_path):
            raise FileNotFoundError(f"历史数据文件不存在: {csv_file_path}")
            
        try:
            # 读取CSV文件
            df = pd.read_csv(csv_file_path)
            
            # 解析draw_numbers列
            numbers_data = []
            for _, row in df.iterrows():
                numbers_str = row['draw_numbers']
                if isinstance(numbers_str, str):
                    # 去除引号并按逗号分割
                    numbers = numbers_str.strip('"').split(',')
                    # 取前5个数字并确保在0-9范围内
                    parsed_numbers = []
                    for num_str in numbers[:5]:
                        try:
                            num = int(num_str.strip())
                            # 使用模运算确保在0-9范围内
                            num = num % 10
                            parsed_numbers.append(num)
                        except ValueError:
                            parsed_numbers.append(0)  # 默认值
                    numbers_data.append(parsed_numbers)
            
            return np.array(numbers_data)
        except Exception as e:
            print(f"加载历史数据失败: {e}")
            return np.array([])
    
    def train(self, history_data=None):
        """
        训练马尔可夫链模型
        
        Args:
            history_data: 历史数据，如果为None则从文件加载
        """
        if history_data is None:
            history_data = self.load_history_data()
            
        if len(history_data) == 0:
            raise ValueError("没有可用的历史数据进行训练")
        
        # 为每个位置构建状态转移矩阵
        for pos in range(5):  # 5个位置
            self._build_transition_matrix(history_data, pos)
        
        # 计算初始概率分布
        self._calculate_initial_probabilities(history_data)
        
        self.is_trained = True
        print(f"✅ 马尔可夫链模型训练完成，使用了{len(history_data)}期历史数据")
    
    def _build_transition_matrix(self, history_data, position):
        """
        为指定位置构建状态转移矩阵
        
        Args:
            history_data: 历史数据
            position: 位置索引(0-4)
        """
        # 初始化转移计数矩阵 (10x10)
        transition_counts = np.zeros((10, 10))
        
        # 统计状态转移
        for i in range(1, len(history_data)):
            prev_state = history_data[i-1][position]
            curr_state = history_data[i][position]
            transition_counts[prev_state][curr_state] += 1
        
        # 转换为概率矩阵（行归一化）
        transition_matrix = np.zeros((10, 10))
        for i in range(10):
            row_sum = np.sum(transition_counts[i])
            if row_sum > 0:
                transition_matrix[i] = transition_counts[i] / row_sum
            else:
                # 如果某状态从未转移，使用均匀分布
                transition_matrix[i] = np.ones(10) / 10
        
        self.transition_matrices[position] = transition_matrix
        
        # 计算位置统计信息
        state_frequencies = np.zeros(10)
        for data in history_data:
            state_frequencies[data[position]] += 1
        state_probabilities = state_frequencies / len(history_data)
        
        self.position_stats[position] = {
            'frequencies': state_frequencies,
            'probabilities': state_probabilities,
            'most_frequent': np.argmax(state_frequencies),
            'least_frequent': np.argmin(state_frequencies)
        }
    
    def _calculate_initial_probabilities(self, history_data):
        """
        计算初始概率分布
        
        Args:
            history_data: 历史数据
        """
        for pos in range(5):
            initial_counts = np.zeros(10)
            for data in history_data:
                initial_counts[data[pos]] += 1
            self.initial_probabilities[pos] = initial_counts / len(history_data)
    
    def predict_next_state(self, current_sequence, position):
        """
        预测指定位置的下一个状态
        
        Args:
            current_sequence: 当前序列
            position: 位置索引
            
        Returns:
            下一个状态的概率分布
        """
        if not self.is_trained:
            raise ValueError("模型未训练，请先调用train()方法")
            
        if position < 0 or position >= 5:
            raise ValueError("位置索引必须在0-4范围内")
        
        current_state = current_sequence[position]
        transition_probs = self.transition_matrices[position][current_state]
        
        return transition_probs
    
    def generate_predictions(self, recent_data, num_predictions=1):
        """
        生成预测结果
        
        Args:
            recent_data: 最近的数据用于初始化
            num_predictions: 预测组数
            
        Returns:
            预测结果列表
        """
        if not self.is_trained:
            raise ValueError("模型未训练，请先调用train()方法")
        
        predictions = []
        
        for _ in range(num_predictions):
            predicted_numbers = []
            
            # 为每个位置生成预测
            for pos in range(5):
                if len(recent_data) > 0:
                    # 使用最近一期数据作为当前状态
                    current_state = recent_data[-1][pos]
                    transition_probs = self.transition_matrices[pos][current_state]
                else:
                    # 使用初始概率分布
                    transition_probs = self.initial_probabilities[pos]
                
                # 根据概率分布随机选择下一个状态
                next_state = np.random.choice(10, p=transition_probs)
                predicted_numbers.append(next_state)
            
            predictions.append({
                'red': predicted_numbers,
                'blue': [],  # 排列5没有蓝球
                'confidence': self._calculate_confidence(predicted_numbers, recent_data)
            })
        
        return predictions
    
    def _calculate_confidence(self, predicted_numbers, recent_data):
        """
        计算预测置信度
        
        Args:
            predicted_numbers: 预测的数字
            recent_data: 最近数据
            
        Returns:
            置信度值(0-1)
        """
        if len(recent_data) == 0:
            return 0.5  # 默认置信度
            
        confidence = 0.0
        for pos in range(5):
            current_state = recent_data[-1][pos]
            predicted_state = predicted_numbers[pos]
            # 转移概率作为置信度的一部分
            transition_prob = self.transition_matrices[pos][current_state][predicted_state]
            confidence += transition_prob
        
        return confidence / 5  # 平均置信度
    
    def get_position_analysis(self, position):
        """
        获取指定位置的分析结果
        
        Args:
            position: 位置索引(0-4)
            
        Returns:
            位置分析字典
        """
        if not self.is_trained:
            raise ValueError("模型未训练，请先调用train()方法")
            
        if position not in self.position_stats:
            return None
            
        stats = self.position_stats[position]
        transition_matrix = self.transition_matrices[position]
        
        # 计算高概率转移
        high_prob_transitions = []
        for i in range(10):
            for j in range(10):
                if transition_matrix[i][j] > 0.2:  # 高于20%的转移概率
                    high_prob_transitions.append((i, j, transition_matrix[i][j]))
        
        return {
            'position': position,
            'state_probabilities': stats['probabilities'],
            'most_frequent_state': stats['most_frequent'],
            'least_frequent_state': stats['least_frequent'],
            'high_probability_transitions': sorted(high_prob_transitions, key=lambda x: x[2], reverse=True)[:5]
        }
    
    def get_model_summary(self):
        """
        获取模型摘要信息
        
        Returns:
            模型摘要字典
        """
        if not self.is_trained:
            return {"status": "未训练"}
            
        summary = {
            "status": "已训练",
            "positions": 5,
            "states_per_position": 10,
            "total_parameters": 5 * 10 * 10,  # 5个位置，每个位置10x10转移矩阵
            "position_details": {}
        }
        
        for pos in range(5):
            if pos in self.position_stats:
                stats = self.position_stats[pos]
                summary["position_details"][pos] = {
                    "most_frequent": int(stats['most_frequent']),
                    "least_frequent": int(stats['least_frequent']),
                    "high_frequency_states": [int(i) for i in np.where(stats['probabilities'] > 0.15)[0]]
                }
        
        return summary


def create_markov_chain_model(csv_file_path):
    """
    创建并训练马尔可夫链模型
    
    Args:
        csv_file_path: 历史数据文件路径
        
    Returns:
        训练好的马尔可夫链模型
    """
    model = PLWMarkovChainModel(csv_file_path)
    model.train()
    return model


# 测试代码
if __name__ == "__main__":
    # 创建测试数据
    test_data = np.array([
        [1, 2, 3, 4, 5],
        [2, 3, 4, 5, 6],
        [3, 4, 5, 6, 7],
        [4, 5, 6, 7, 8],
        [5, 6, 7, 8, 9],
        [6, 7, 8, 9, 0],
        [7, 8, 9, 0, 1],
        [8, 9, 0, 1, 2],
        [9, 0, 1, 2, 3],
        [0, 1, 2, 3, 4]
    ])
    
    # 创建并训练模型
    model = PLWMarkovChainModel()
    model.train(test_data)
    
    # 生成预测
    predictions = model.generate_predictions(test_data[-3:], 3)
    print("_predictions:")
    for i, pred in enumerate(predictions):
        print(f"  预测 {i+1}: {pred['red']} (置信度: {pred['confidence']:.3f})")
    
    # 位置分析
    print("\n位置分析:")
    for pos in range(5):
        analysis = model.get_position_analysis(pos)
        if analysis:
            print(f"  位置 {pos}: 最频繁数字={analysis['most_frequent_state']}, "
                  f"高概率转移={len(analysis['high_probability_transitions'])}个")